package org.zhang.core.asm;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;
import org.zhang.core.config.ProfilingFilter;

import java.io.File;
import java.io.IOException;
import java.io.OutputStream;
import java.lang.instrument.ClassFileTransformer;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.ProtectionDomain;

public class MyClassFileTransformer implements ClassFileTransformer {
    public MyClassFileTransformer(String args) {
        System.out.println("初始化参数：" + args);
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {
        try {
            if (ProfilingFilter.isNotNeedInject(className)) {
                return classfileBuffer;
            }
            byte[] transformedBytes = transformClass(classfileBuffer, className, loader);
//            outputClassFile(transformedBytes, className);
            return transformedBytes;
        } catch (Throwable e) {
            handleTransformationError(e, className);
            return classfileBuffer;
        }
    }

    private byte[] transformClass(byte[] classfileBuffer, String className, ClassLoader loader) {
        ClassReader cr = new ClassReader(classfileBuffer);
        ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_FRAMES) {
            @Override
            protected ClassLoader getClassLoader() {
                return loader;
            }
        };
        ClassVisitor cv = new MethodInterceptorClassVisitor(Opcodes.ASM7, cw, className);
        cr.accept(cv, ClassReader.EXPAND_FRAMES);
        return cw.toByteArray();
    }

    /**
     * 输出插桩后的字节码文件
     *
     * @param bytes
     * @param className
     */
    private void outputClassFile(byte[] bytes, String className) {
        try (OutputStream fos = Files.newOutputStream(Paths.get(getOutputFilePath(className)))) {
            fos.write(bytes);
            System.out.println("ASM 类输出路径：" + getOutputFilePath(className));
        } catch (IOException e) {
            handleOutputError(e, className);
        }
    }

    /**
     * 根据需要更改文件路径生成逻辑。
     */
    private String getOutputFilePath(String className) throws IOException {
        // 定义基础路径
//        String basePath = "C:\\Users\\32760\\Desktop\\classes";

        // 将类名中的点替换为文件分隔符
        String relativePath = "asm-classes" + File.separatorChar + className
                .replace('.', File.separatorChar)
                .replace('/', File.separatorChar)
                .replace('\\', File.separatorChar)
                + "SQM.class";

        // 完整的文件路径
//        String fullPath = basePath + File.separator + relativePath;
        String fullPath = relativePath;

        // 检查目录是否存在，如果不存在则创建
        Files.createDirectories(Paths.get(fullPath).getParent());
        return fullPath;
    }

    private void handleTransformationError(Throwable e, String className) {
        System.err.println("转换类 " + className + " 时发生错误：" + e.getMessage());
        // 添加更多的错误处理逻辑，例如记录到日志
        e.printStackTrace();
    }

    private void handleOutputError(IOException e, String className) {
        System.err.println("写入类文件 " + className + " 时发生错误：" + e.getMessage());
        // 添加更多的错误处理逻辑，例如记录到日志
    }
}